import cv2
import numpy as np
import os

DST_IMG_DIR = "./dst_img/"
DST_SPLIT_IMG_DIR = "./dst_split_img/"
LETTER_TEMPLATE_DIR = "./letter_template/"
LETTER_TEMPLATE_DEFAULT_DIR = "./letter_template/"
LETTER_TEST_DIR = "./letter_test/"
SRC_IMG_DIR = "./src_img/"
TEST_SRC_IMG_DIR = "./test_src_img/"

SPLIT_WIDTH = 40
SPLIT_HEIGHT = 40

MAX_WIDTH = 0
MAX_HEIGHT = 0

IS_SAVE_FILE = True

def setSaveFileMode(flag):
    global IS_SAVE_FILE
    IS_SAVE_FILE = flag
    pass

def setTestMode(flag):
    global LETTER_TEMPLATE_DIR
    if flag:
        LETTER_TEMPLATE_DIR = LETTER_TEST_DIR
    else:
        LETTER_TEMPLATE_DIR = LETTER_TEMPLATE_DEFAULT_DIR
    pass

def show(name,img):
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    pass

def get_file_name(file_path):
    dir_index = file_path.rindex(os.sep) + 1
    file_name = file_path[dir_index:].split(".")[0]
    return file_name

def read_image(file_path):
    img = cv2.imread(file_path)
    return img

def pre_process_image(img, file_name):
    # global IS_SAVE_FILE

    # 去除边缘
    img = img[2:-2, 2:-2]
    # print(img.shape)
    # show("img", img)

    #得到灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # show("gray", gray)

    #去除噪音
    blur = cv2.medianBlur(gray, 3)
    # show("blur", blur)

    temp = gray.mean().item()
    #二值化
    ret, threshold = cv2.threshold(blur, temp, 255, cv2.THRESH_BINARY)
    # show("threshold", threshold)

    #保存二值化图片
    if IS_SAVE_FILE:
        cv2.imwrite(DST_IMG_DIR + file_name + "_threshold.png", threshold)
    return threshold


def find_counters(img):
    # v4.x
    # contours, hierarchy= cv2.findContours(threshold, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    # v3.x
    binary, contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    # print(type(contours))
    # print("contours length:",len(contours))
    return contours


def get_filter_rect(contours,img,file_name):
    rect = img
    list = []
    # 边界矩形面积
    s = img.shape[0] * img.shape[1]
    min_x = img.shape[1]
    max_x = 0
    min_y = img.shape[0]
    max_y = 0
    middle_min_y = 10

    index_list = {}

    for index,cnt in enumerate(contours):
        x, y, w, h = cv2.boundingRect(cnt)
        # extent = float(cv2.contourArea(cnt)) / (w * h)
        if(x > 0):
            min_x = min(min_x,x)
        max_x = max(max_x,x + w)
        if (y > 0):
            min_y = min(min_y,y)
        max_y = max(max_y,y + h)

        # if w*h < s/20.0 or w*h > s/3.0: continue
        if w*h < 360 or w*h > 1296: continue

        middle_min_y = (middle_min_y + y) / 2
        letter_index = round(x / (img.shape[1] / 4))
        index_list[letter_index] = [x, y, w, h]
        list.append([x, y, w, h])
        # print("轮廓面积与边界矩形比：", extent, w * h)
        rect = cv2.rectangle(rect, (x+2, y+2), (x + w, y + h), (0, 0, 255), 2)
        # show("rect", rect)

    for i in range(4):
        # last_y = min_y

        #取上一个字符的y
        # if(i > 0 and index_list.get(i - 1) is not None):
        #     last_y = index_list.get(i - 1)[1]

        if index_list.get(i) is None:
            # d = int((max_x - min_x) / 4)
            d = 30
            x = d * i + 2
            if i == 0:
                x = 6
            y = int(middle_min_y)
            list.append([x, y,  d, SPLIT_HEIGHT])
            rect = cv2.rectangle(rect, (x+2, y+2), (x + d, y + SPLIT_HEIGHT), (255, 0, 0), 2)

    result_rect = (min_x, min_y, max_x, max_y)
    # rect = cv2.rectangle(rect, (min_x, min_y), (max_x, max_y), (255, 0, 0), 3)
    # show("rect", rect)

    # print("min x-max x,min y-max y",min_x, max_x, min_y, max_y)
    if IS_SAVE_FILE:
        cv2.imwrite(DST_IMG_DIR + file_name + "_threshold_rect.png", rect)
    return list, result_rect


def take_first(elem):
    return elem[0]

def fill_rect(rect,fill_width,fill_height):
    w = fill_width - rect.shape[1]
    p_left = int(w / 2)
    p_right = w - p_left

    h = fill_height - rect.shape[0]
    p_top = int(h / 2)
    p_bottom = h - p_top

    if rect.shape[1] < fill_width:
        rect = np.pad(rect, ((0, 0), (p_left, p_right)), 'constant', constant_values=(255, 255))
    elif rect.shape[1] > fill_width:
        if p_right >= 0:
            p_right = rect.shape[1]
        rect = rect[:, abs(p_left): p_right ]

    if rect.shape[0] < fill_height:
        rect = np.pad(rect, ((p_top, p_bottom), (0, 0)), 'constant', constant_values=(255, 255))
    elif rect.shape[0] > fill_height:
        # print("wh",w,h,p_top,p_bottom,rect.shape)
        if p_bottom >= 0:
            p_bottom = rect.shape[0]
        rect = rect[abs(p_top): p_bottom,:]

    return rect

def split_rect_img(file_path,img,rect_list,result_rect):
    global MAX_WIDTH
    global MAX_HEIGHT
    if(len(rect_list) == 0):
        return

    file_name = get_file_name(file_path)
    full_split = len(rect_list) == 4
    result_list = []
    #按x大小排序
    rect_list.sort(key=take_first)

    for index,rect in enumerate(rect_list):

        x = max(0,int(rect[0]) - 2)
        y = max(0,int(rect[1]) - 2)
        w = int(rect[2]) + 4
        h = int(rect[3]) + 4

        img_rect = img[y:y + h, x:x + w]
        # print("img_rect shape", img_rect.shape)
        MAX_WIDTH = max(MAX_WIDTH, w)
        MAX_HEIGHT = max(MAX_HEIGHT, h)

        # show("img_rect", img_rect)
        img_rect = fill_rect(img_rect, SPLIT_WIDTH, SPLIT_HEIGHT)
        if(img_rect.shape[0] != SPLIT_HEIGHT or img_rect.shape[1] != SPLIT_WIDTH):
            print("invalid shape",img_rect.shape)

        if IS_SAVE_FILE :
            file_dir = DST_SPLIT_IMG_DIR + file_name
            if not os.path.exists(file_dir):
                os.makedirs(file_dir)
            dst_file_name = file_dir + "/" + str(index) + ".png"
            cv2.imwrite(dst_file_name, img_rect)

        if full_split and IS_SAVE_FILE:
            letter_dir = LETTER_TEMPLATE_DIR + file_name[index]
            if not os.path.exists(letter_dir):
                os.makedirs(letter_dir)
            dst_file_name = letter_dir + "/" + file_name + "_" + str(index) + ".png"
            cv2.imwrite(dst_file_name, img_rect)

        result_list.append(img_rect)
    return result_list

def split_image(file_path):
    file_name = get_file_name(file_path)
    img = read_image(file_path)

    #验证码预处理
    threshold = pre_process_image(img, file_name)

    #查找轮廓边界列表
    contours = find_counters(threshold)

    #过滤合适的轮廓矩形列表
    rect_list,result_rect = get_filter_rect(contours, img, file_name)

    #分割矩形图片
    return split_rect_img(file_path, threshold, rect_list, result_rect)

def split_image_dir(dir):
    if not os.path.exists(DST_IMG_DIR):
        os.makedirs(DST_IMG_DIR)

    for root, dirs, files in os.walk(dir):
        for file in files:
            # 获取文件路径
            file_path = os.path.join(root, file)
            # print(file_path)
            split_image(file_path)
    pass

def split_test_image():
    setTestMode(True)
    split_image_dir(TEST_SRC_IMG_DIR)
    # split_image(TEST_SRC_IMG_DIR + "BHJ3.jpg")
    # split_image(TEST_SRC_IMG_DIR + "Z2HS.jpg")
    pass

if __name__ == '__main__':
    # split_image_dir(SRC_IMG_DIR)

    split_test_image()

    print(MAX_WIDTH)
    print(MAX_HEIGHT)
    pass



